# coding: utf-8
# 2021/6/21 @ tongshiwei
import numpy as np
import torch
from torch import nn

from .base import Module


class DINANet(Module):
    def __init__(
        self,
        user_num,
        item_num,
        hidden_dim,
        max_slip=0.4,
        max_guess=0.4,
        *args,
        **kwargs
    ):
        super(DINANet, self).__init__()
        self._user_num = user_num
        self._item_num = item_num
        self.step = 0
        self.max_step = 1000
        self.max_slip = max_slip
        self.max_guess = max_guess

        self.guess = nn.Embedding(self._item_num, 1)
        self.slip = nn.Embedding(self._item_num, 1)
        self.theta = nn.Embedding(self._user_num, hidden_dim)

    def forward(self, user, item, knowledge, *args):
        theta = self.theta(user)
        slip = torch.squeeze(torch.sigmoid(self.slip(item)) * self.max_slip)
        guess = torch.squeeze(torch.sigmoid(self.guess(item)) * self.max_guess)
        if self.training:
            n = torch.sum(knowledge * (torch.sigmoid(theta) - 0.5), dim=1)
            t, self.step = max(
                (np.sin(2 * np.pi * self.step / self.max_step) + 1) / 2 * 100, 1e-6
            ), (self.step + 1 if self.step < self.max_step else 0)
            return torch.sum(
                torch.stack([1 - slip, guess]).T
                * torch.softmax(torch.stack([n, torch.zeros_like(n)]).T / t, dim=-1),
                dim=1,
            )
        else:
            n = torch.prod(knowledge * (theta >= 0) + (1 - knowledge), dim=1)
            return (1 - slip) ** n * guess ** (1 - n)
